{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://pymt.readthedocs.io\"><img style=\"float: left\" src=\"../../media/pymt-logo-header-text.png\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Couple two models in *pymt*\n",
    "\n",
    "In this exercise we will couple two component using *pymt*. We will couple the Coastline Evolution Model (CEM),\n",
    "which models the transport of sediment alongshore, with a wave model that provides incoming wave direction and\n",
    "characteristics.\n",
    "\n",
    "* Explore the base-case river simulation\n",
    "* How does a river system respond to climate change?\n",
    "* How do humans affect river sediment loads?\n",
    "\n",
    "## The Coastline Evolution Model (CEM)\n",
    "\n",
    "The Coastline Evolution Model (CEM) addresses predominately sandy, wave-dominated coastlines on time-scales ranging from years to millenia and on spatial scales ranging from kilometers to hundreds of kilometers. Shoreline evolution results from gradients in wave-driven alongshore sediment transport.\n",
    "\n",
    "At its most basic level, the model follows the standard *one-line* modeling approach, where the cross-shore dimension is collapsed into a single data point. However, the model allows the planview shoreline to take on arbitrary local orientations, and even fold back upon itself, as complex shapes such as capes and spits form under some wave climates (distributions of wave influences from different approach angles). The model works on a 2D grid.\n",
    "\n",
    "CEM has been used to represent varying geology underlying a sandy coastline and shoreface in a simplified manner and enables the simulation of coastline evolution when sediment supply from an eroding shoreface may be constrained. CEM also supports the simulation of human manipulations to coastline evolution through beach nourishment or hard structures.\n",
    "\n",
    "CEM authors & developers include:\n",
    "* Andrew Ashton\n",
    "* Brad Murray\n",
    "* Jordan Slot\n",
    "* Jaap Nienhuis and others.\n",
    "\n",
    "This version is adapted from a CSDMS teaching notebook, listed below. \n",
    "It has been created by Irina Overeem, October 2019 for a Sedimentary Modeling course.\n",
    "\n",
    "### Key References\n",
    "\n",
    "Ashton, A.D., Murray, B., Arnault, O. 2001. Formation of coastline features by large-scale instabilities induced by high-angle waves, Nature 414.\n",
    "\n",
    "Ashton, A. D., and A. B. Murray (2006), High-angle wave instability and emergent shoreline shapes: 1. Modeling of sand waves, flying spits, and capes, J. Geophys. Res., 111, F04011, doi:10.1029/2005JF000422.\n",
    "\n",
    "\n",
    "### Links\n",
    "* [CEM source code](https://github.com/csdms/cem-old/tree/mcflugen/add-function-pointers): Look at the files that have *deltas* in their name.\n",
    "* [CEM description on CSDMS](http://csdms.colorado.edu/wiki/Model_help:CEM): Detailed information on the CEM model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interacting with the Coastline Evolution Model BMI using Python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the `Cem` model into your environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymt.models\n",
    "cem = pymt.models.Cem()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even though we can't run our model yet, we can still get some information about it. Some things we can do with our model are to get help, and get the names of the input variables or output variables. In looking at the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cem.input_var_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cem.output_var_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also get information about specific variables. Here we'll look at some info about wave direction. This is the main input of the CEM model. What do you think the more conventional names for these variables are?\n",
    "\n",
    "| Conventional Name      | Standard Name                                                       |\n",
    "| :--------------------- | :------------------------------------------------------------------ |\n",
    "| ???                    | sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity |\n",
    "| ???                    | sea_surface_water_wave__period                                      |\n",
    "| ???                    | sea_surface_water_wave__height                                      |\n",
    "\n",
    "To help us out, we can get some additional information about each of the variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "angle_name = 'sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity'\n",
    "\n",
    "print(\"Data type: %s\" % cem.var_type(angle_name))\n",
    "print(\"Units: %s\" % cem.var_units(angle_name))\n",
    "print(\"Grid id: %d\" % cem.var_grid(angle_name))\n",
    "print(\"Number of elements in grid: %d\" % cem.grid_node_count(0))\n",
    "print(\"Type of grid: %s\" % cem.grid_type(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now get the model ready for time stepping. Remember the lifecycle of the model is:\n",
    "* *setup*\n",
    "* *initialize*\n",
    "* *update*\n",
    "* *finalize*\n",
    "\n",
    "For this example we'll set up a simulation with a grid of 100 rows and 200 columns with a grid\n",
    "resolution of 200.0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = cem.setup(number_of_rows=100, number_of_cols=200, grid_spacing=200.)\n",
    "cem.initialize(*args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing a model's state\n",
    "\n",
    "Because the CEM model has input variables (unlike *HydroTrend* in the previous example), we\n",
    "are able to change variables within the CEM model. The is done with the ***set_value***\n",
    "method.\n",
    "\n",
    "For our first example we'll set the incoming wave height, period, and angle (in radians)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cem.set_value(\"sea_surface_water_wave__height\", 1.5)\n",
    "cem.set_value(\"sea_surface_water_wave__period\", 7.)\n",
    "cem.set_value(\"sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity\", 0.0 * np.pi / 180.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grids\n",
    "\n",
    "This example is also different from the previous example in that it generates output that\n",
    "is on a grid (as opposed to scalar data). The main output for CEM is *sea_water__depth*\n",
    "and operates on a grid whose size we set when we *setup* this simulation.\n",
    "\n",
    "*pymt* models can have multiple grids. This allows for models to calculate some of\n",
    "its state variables on scalars and others of 2D grids, for example. Models could also\n",
    "maintain several grids of differing resolution. In *pymt* each grid has an ID\n",
    "associated with it.\n",
    "\n",
    "We use the ***var_grid*** function to get the grid ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cem.var_grid(\"sea_water__depth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have a grid ID, we can then use the *pymt* *grid_* methods to get additional information\n",
    "about each of the grids. Because this grid is uniform rectilinear (as returned by the\n",
    "***grid_type*** method below), it is described by a set of  methods that are only available\n",
    "for grids of this type. These methods include:\n",
    "* get_grid_shape\n",
    "* get_grid_spacing\n",
    "* get_grid_origin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Grid type: {0}\".format(cem.grid_type(2)))\n",
    "print(\"Grid rank: {0}\".format(cem.grid_ndim(2)))\n",
    "print(\"Grid shape: {0}\".format(cem.grid_shape(2)))\n",
    "print(\"Grid spacing: {0}\".format(cem.grid_spacing(2)))\n",
    "print(\"Grid origin: {0}\".format(cem.grid_origin(2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we know a little more about the grid, let's plot it with the current values of\n",
    "water depth.\n",
    "\n",
    "Here I define a convenience function for plotting the water depth and making it look\n",
    "pretty. You don't need to worry too much about it's internals for this tutorial.\n",
    "It just saves typing later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_coast(spacing, z):\n",
    "    import matplotlib.pyplot as plt\n",
    "\n",
    "    xmin, xmax = 0., z.shape[1] * spacing[0] * 1e-3\n",
    "    ymin, ymax = 0., z.shape[0] * spacing[1] * 1e-3\n",
    "\n",
    "    plt.imshow(z, extent=[xmin, xmax, ymin, ymax], origin=\"lower\", cmap=\"ocean\")\n",
    "    plt.colorbar().ax.set_ylabel(\"Water Depth (m)\")\n",
    "    plt.xlabel(\"Along shore (km)\")\n",
    "    plt.ylabel(\"Cross shore (km)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.empty(cem.grid_shape(2), dtype=float)\n",
    "\n",
    "cem.get_value(\"sea_water__depth\", out=z)\n",
    "plot_coast(cem.grid_spacing(2), z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have alredy set the incoming wave characteristics, but we've yet to add any sediment to\n",
    "the system.\n",
    "\n",
    "From the list of input variables, can you tell which one might be the one we're looking\n",
    "for?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in cem.input_var_names:\n",
    "    print(\"{0:70s} [{1}]\".format(name, cem.var_units(name)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The one we want is *land_surface_water_sediment~bedload__mass_flow_rate*. Now have a look\n",
    "at what sort of grid its defined on and how we could change its value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cem.var_grid(\"land_surface_water_sediment~bedload__mass_flow_rate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that it's on the same grid as water depth. To add sediment, we need to:\n",
    "1. allocate an array to hold sediment discharge values\n",
    "2. set values of the sediment discharge array\n",
    "3. pass this new sediment discharge array into CEM\n",
    "\n",
    "I've placed the sediment discharge in the horizontal center of the grid (column 100 of 200) and\n",
    "along the bottom. The sediment will be routed in a straight line until it hits the coast.\n",
    "\n",
    "You don't need to do this, though. Feel free to add sediment in another location (or even multiple\n",
    "locations!) or change the amount of sediment. Note that the CEM model is sensitive to the balance of\n",
    "wave energy to sediment input. If you go too far from the defaults, you may get some \"interesting\"\n",
    "results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qs = np.zeros_like(z)\n",
    "qs[0, 100] = 750\n",
    "cem.set_value(\"land_surface_water_sediment~bedload__mass_flow_rate\", qs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's time step through the model. For every iteration of the for-loop we set the sediment\n",
    "discharge (***set_value***), and then update the model to the next time (***update_until***)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for time in tqdm(range(3000)):\n",
    "    cem.set_value('land_surface_water_sediment~bedload__mass_flow_rate', qs)\n",
    "    cem.update_until(time)\n",
    "\n",
    "cem.get_value('sea_water__depth', out=z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look to see what sort of delta we've created after 3000.0 days."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cem.get_value('sea_water__depth', out=z)\n",
    "plot_coast(cem.grid_spacing(2), z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "Now play with the CEM model on your own. To make things easier, I've placed all of the steps\n",
    "to run CEM into a single cell below. Please feel free to modify!\n",
    "\n",
    "Some ideas\n",
    "1. Modify wave energy vs sediment load\n",
    "1. Change the incoming wave angle\n",
    "1. Modify the river position so that it moves with time\n",
    "1. Pick wave height and period from some probability density function\n",
    "1. Add a second, or third, or fourth river\n",
    "1. Increase the sediment load or have it vary with time\n",
    "1. Make a movie of the evolving delta\n",
    "\n",
    "Anything else?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymt.models\n",
    "cem = pymt.models.Cem()\n",
    "\n",
    "args = cem.setup(number_of_rows=100, number_of_cols=200, grid_spacing=200.)\n",
    "cem.initialize(*args)\n",
    "\n",
    "qs = np.zeros(cem.grid_shape(2), dtype=float)\n",
    "qs[0, 100] = 750\n",
    "\n",
    "for time in tqdm(range(3000)):\n",
    "    cem.set_value(\"sea_surface_water_wave__height\", 1.5)\n",
    "    cem.set_value(\"sea_surface_water_wave__period\", 7.)\n",
    "    cem.set_value(\n",
    "        \"sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity\",\n",
    "        0. * np.pi / 180.,\n",
    "    )\n",
    "\n",
    "    cem.set_value('land_surface_water_sediment~bedload__mass_flow_rate', qs)\n",
    "    cem.update_until(time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.empty(cem.grid_shape(2), dtype=float)\n",
    "cem.get_value('sea_water__depth', out=z)\n",
    "\n",
    "plot_coast(cem.grid_spacing(2), z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def avulse_river(river_x, stddev=1.0, x_min=None, x_max=None):\n",
    "    river_x += np.random.normal(0.0, stddev)\n",
    "    if x_max is not None and river_x >= 200:\n",
    "        river_x = river_x - 200\n",
    "    if x_min is not None and river_x < 0:\n",
    "        river_x = 200 + river_x\n",
    "    return river_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "river_x = 100.0\n",
    "river_i = []\n",
    "for time in range(3000):\n",
    "    river_x = avulse_river(river_x, x_min=0.0, x_max=200.0)\n",
    "    river_i.append(int(river_x))\n",
    "plt.plot(river_i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "### Couple models\n",
    "\n",
    "Instead of using constant wave characteristics, let's now couple the CEM component with a wave\n",
    "gererator component.\n",
    "\n",
    "### Waves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymt.models import Waves\n",
    "\n",
    "waves = Waves()\n",
    "args = waves.setup(angle_asymmetry=0.2, angle_highness_factor=0.5)\n",
    "waves.initialize(*args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "angles = np.zeros(1000)\n",
    "for day in range(1000):\n",
    "    waves.update()\n",
    "    angles[day] = waves.get_value(\"sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.hist(angles * 180.0 / np.pi, bins=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now add the waves component to our coupling script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pymt.models\n",
    "\n",
    "waves = pymt.models.Waves()\n",
    "args = waves.setup(angle_asymmetry=0.2, angle_highness_factor=0.5)\n",
    "waves.initialize(*args)\n",
    "\n",
    "cem = pymt.models.Cem()\n",
    "args = cem.setup(number_of_rows=100, number_of_cols=200, grid_spacing=200.)\n",
    "cem.initialize(*args)\n",
    "\n",
    "qs = np.zeros(cem.grid_shape(2), dtype=float)\n",
    "qs[0, 100] = 500\n",
    "\n",
    "for time in tqdm(range(3000)):\n",
    "    waves.update()\n",
    "    angle = waves.get_value(\"sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity\")\n",
    "    \n",
    "    cem.set_value(\"sea_surface_water_wave__height\", 1.5)\n",
    "    cem.set_value(\"sea_surface_water_wave__period\", 7.0)\n",
    "    cem.set_value(\n",
    "        \"sea_surface_water_wave__azimuth_angle_of_opposite_of_phase_velocity\",\n",
    "        angle,\n",
    "    )\n",
    "\n",
    "    cem.set_value('land_surface_water_sediment~bedload__mass_flow_rate', qs)\n",
    "    cem.update_until(time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = cem.get_value('sea_water__depth').reshape(cem.grid_shape(2))\n",
    "plot_coast(cem.grid_spacing(2), z)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
